import json
import time
import threading
from hqyj_mqtt import Mqtt_Clt
import base64
import numpy as np
import cv2
from predict import predict_resnet

# 定义光推杆序号常量
load_rod = "first"
ripe_rod = "second"
half_ripe_rod = "third"
raw_rod = "fourth"
all_rod = "all"

# 定义光电对管状态字典
ripe_switch = {"first_switch": False}
half_ripe_switch = {"second_switch": False}
raw_switch = {"third_switch": False}

# 加载推杆控制线程引用（用于可能的停止控制）
load_rod_thread = None
stop_load_rod = False


def push_pull(rod):
    """单次推出收回控制"""
    mqtt_client.send_json_msg(json.dumps({"rod_control": f"{rod}_push"}))
    time.sleep(0.5)
    mqtt_client.send_json_msg(json.dumps({"rod_control": f"{rod}_pull"}))


def start_continuous_load_rod(interval=0.8):
    """启动加载推杆的持续循环"""
    global stop_load_rod

    def worker():
        while not stop_load_rod:
            # 推出
            mqtt_client.send_json_msg(json.dumps({"rod_control": f"{load_rod}_push"}))
            time.sleep(0.5)  # 推出保持时间
            # 收回
            mqtt_client.send_json_msg(json.dumps({"rod_control": f"{load_rod}_pull"}))
            time.sleep(interval - 0.5)  # 间隔时间

    stop_load_rod = False
    thread = threading.Thread(target=worker, daemon=True)
    thread.start()
    return thread


def stop_continuous_load_rod():
    """停止加载推杆的持续循环"""
    global stop_load_rod
    stop_load_rod = True


if __name__ == '__main__':
    mqtt_client = Mqtt_Clt("127.0.0.1", 21883, "bb", "aa", 60)
    print("开始控制")
    mqtt_client.control_device("conveyor", 'run')

    # 控制所有推杆收回初始位置
    mqtt_client.control_device("rod_control", "all_pull")

    # 启动加载推杆的持续循环（2秒一个周期）
    load_rod_thread = start_continuous_load_rod(interval=0.8)

    result = None
    while True:
        msg = mqtt_client.mqtt_queue.get()

        if 'image' in msg:
            image_data = base64.b64decode(msg['image'])
            image_array = np.frombuffer(image_data, np.uint8)
            image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
            result = predict_resnet(image)
            print(f"result: {result}")

        else:

            if (result == "ripe") and (ripe_switch == msg):
                threading.Timer(0.0, push_pull, args=(ripe_rod,)).start()
            elif (result == "half-ripe") and (half_ripe_switch == msg):
                threading.Timer(0.0, push_pull, args=(half_ripe_rod,)).start()
            elif (result == "raw") and (raw_switch == msg):
                threading.Timer(0.0, push_pull, args=(raw_rod,)).start()